"""SSE (Server-Sent Events) 处理器."""

import asyncio
import json
import logging
from datetime import datetime
from typing import AsyncGenerator, Optional, Any, Dict

from fastapi import Request
from fastapi.responses import StreamingResponse
from sse_starlette.sse import EventSourceResponse

logger = logging.getLogger(__name__)


class SSEHandler:
    """SSE事件处理器."""
    
    def __init__(self):
        self.active_streams: Dict[str, bool] = {}
        self.event_id_counter = 0
        
    def create_event(
        self,
        data: Any,
        event: Optional[str] = None,
        event_id: Optional[str] = None,
        retry: Optional[int] = None
    ) -> str:
        """创建SSE事件字符串.
        
        Args:
            data: 事件数据
            event: 事件类型
            event_id: 事件ID
            retry: 重试间隔（毫秒）
            
        Returns:
            SSE事件字符串
        """
        message_parts = []
        
        # 添加事件ID
        if event_id is None:
            self.event_id_counter += 1
            event_id = str(self.event_id_counter)
        message_parts.append(f"id: {event_id}")
        
        # 添加事件类型
        if event:
            message_parts.append(f"event: {event}")
            
        # 添加重试间隔
        if retry:
            message_parts.append(f"retry: {retry}")
            
        # 添加数据
        if isinstance(data, dict):
            data_str = json.dumps(data, ensure_ascii=False)
        else:
            data_str = str(data)
            
        # SSE支持多行数据
        for line in data_str.split('\n'):
            message_parts.append(f"data: {line}")
            
        # SSE事件以双换行结束
        return '\n'.join(message_parts) + '\n\n'
        
    async def create_stream(
        self,
        task_id: str,
        request: Request,
        generator_func: AsyncGenerator
    ) -> EventSourceResponse:
        """创建SSE流响应.
        
        Args:
            task_id: 任务ID
            request: FastAPI请求对象
            generator_func: 异步生成器函数
            
        Returns:
            SSE响应对象
        """
        async def event_generator():
            """内部事件生成器."""
            self.active_streams[task_id] = True
            
            try:
                # 发送初始连接事件
                yield {
                    "event": "connection",
                    "data": {
                        "task_id": task_id,
                        "timestamp": datetime.utcnow().isoformat()
                    }
                }
                
                # 转发生成器的内容
                async for item in generator_func:
                    # 检查客户端是否断开
                    if await request.is_disconnected():
                        logger.info(f"SSE客户端断开: task_id={task_id}")
                        break
                        
                    # 检查流是否被取消
                    if not self.active_streams.get(task_id, False):
                        logger.info(f"SSE流被取消: task_id={task_id}")
                        break
                        
                    yield item
                    
                # 发送完成事件
                yield {
                    "event": "complete",
                    "data": {
                        "task_id": task_id,
                        "timestamp": datetime.utcnow().isoformat()
                    }
                }
                
            except asyncio.CancelledError:
                logger.info(f"SSE流被取消: task_id={task_id}")
                raise
            except Exception as e:
                logger.error(f"SSE流错误: task_id={task_id}, error={e}")
                yield {
                    "event": "error",
                    "data": {
                        "task_id": task_id,
                        "error": str(e),
                        "timestamp": datetime.utcnow().isoformat()
                    }
                }
            finally:
                # 清理
                if task_id in self.active_streams:
                    del self.active_streams[task_id]
                    
        return EventSourceResponse(
            event_generator(),
            media_type="text/event-stream",
            headers={
                "Cache-Control": "no-cache",
                "Connection": "keep-alive",
                "X-Accel-Buffering": "no"  # 禁用nginx缓冲
            }
        )
        
    async def send_token_stream(
        self,
        task_id: str,
        token_generator: AsyncGenerator[str, None]
    ) -> AsyncGenerator[dict, None]:
        """发送token流.
        
        Args:
            task_id: 任务ID
            token_generator: token生成器
            
        Yields:
            SSE事件字典
        """
        token_count = 0
        buffer = ""
        
        async for token in token_generator:
            token_count += 1
            buffer += token
            
            # 每收集一定数量的token或遇到标点符号时发送
            if len(buffer) >= 10 or token in '.!?。！？\n':
                yield {
                    "event": "token",
                    "data": {
                        "content": buffer,
                        "token_count": token_count,
                        "timestamp": datetime.utcnow().isoformat()
                    }
                }
                buffer = ""
                
            # 添加流量控制
            await asyncio.sleep(0.01)
            
        # 发送剩余的buffer
        if buffer:
            yield {
                "event": "token",
                "data": {
                    "content": buffer,
                    "token_count": token_count,
                    "timestamp": datetime.utcnow().isoformat()
                }
            }
            
    def cancel_stream(self, task_id: str) -> bool:
        """取消SSE流.
        
        Args:
            task_id: 任务ID
            
        Returns:
            是否成功取消
        """
        if task_id in self.active_streams:
            self.active_streams[task_id] = False
            logger.info(f"SSE流已取消: task_id={task_id}")
            return True
        return False
        
    def is_active(self, task_id: str) -> bool:
        """检查流是否活跃.
        
        Args:
            task_id: 任务ID
            
        Returns:
            流是否活跃
        """
        return self.active_streams.get(task_id, False)
        
    def get_active_streams(self) -> list:
        """获取活跃流列表.
        
        Returns:
            活跃的任务ID列表
        """
        return list(self.active_streams.keys())
        

# 全局SSE处理器实例
sse_handler = SSEHandler()